package org.ruoyi.interview.config;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.WebSocket;
import java.net.http.WebSocket.Listener;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

public class XunfeiWebSocketClient implements Listener {

    private final StringBuilder resultBuilder = new StringBuilder();
    private final CountDownLatch closeLatch = new CountDownLatch(1);
    private volatile boolean isCompleted = false;
    private final String url;
    private final String appId;
    private WebSocket webSocket;

    public XunfeiWebSocketClient(String url, String appId) {
        this.url = url;
        this.appId = appId;
    }

    /**
     * 建立连接
     */
    public void connectBlocking(String authUrl) throws InterruptedException {
        HttpClient client = HttpClient.newHttpClient();
        CompletableFuture<WebSocket> ws = client.newWebSocketBuilder()
                .connectTimeout(Duration.ofSeconds(10))
                .buildAsync(URI.create(url), this);

        webSocket = ws.join(); // 等待连接建立
    }

    /**
     * 发送文本消息
     */
    public void sendText(String text) {
        if (webSocket != null) {
            webSocket.sendText(text, true);
        }
    }

    /**
     * 关闭连接
     */
    public void closeBlocking() throws InterruptedException {
        if (webSocket != null) {
            webSocket.sendClose(WebSocket.NORMAL_CLOSURE, "ok").join();
        }
        // 等待连接关闭
        closeLatch.await(5, TimeUnit.SECONDS);
    }

    @Override
    public void onOpen(WebSocket webSocket) {
        System.out.println("WebSocket连接已建立");
        webSocket.request(1);
    }

    @Override
    public CompletionStage<?> onText(WebSocket webSocket, CharSequence data, boolean last) {
        try {
            String message = data.toString();
            // 解析返回的JSON消息
            cn.hutool.json.JSONObject jsonObject = new cn.hutool.json.JSONObject(message);
            cn.hutool.json.JSONObject header = jsonObject.getJSONObject("header");
            cn.hutool.json.JSONObject payload = jsonObject.getJSONObject("payload");

            if (payload != null && payload.containsKey("choices")) {
                cn.hutool.json.JSONObject choices = payload.getJSONObject("choices");
                if (choices.containsKey("text")) {
                    cn.hutool.json.JSONArray textArray = choices.getJSONArray("text");
                    if (textArray != null && !textArray.isEmpty()) {
                        cn.hutool.json.JSONObject textObj = textArray.getJSONObject(0);
                        String content = textObj.getStr("content");
                        if (content != null) {
                            resultBuilder.append(content);
                        }
                    }
                }
            }

            // 检查是否结束
            if (header != null && header.getInt("code") == 0) {
                if (header.containsKey("status") && header.getInt("status") == 2) {
                    isCompleted = true;
                    closeLatch.countDown();
                }
            }

        } catch (Exception e) {
            e.printStackTrace();
        }

        webSocket.request(1);
        return Listener.super.onText(webSocket, data, last);
    }

    @Override
    public CompletionStage<?> onBinary(WebSocket webSocket, ByteBuffer data, boolean last) {
        webSocket.request(1);
        return Listener.super.onBinary(webSocket, data, last);
    }

    @Override
    public CompletionStage<?> onClose(WebSocket webSocket, int statusCode, String reason) {
        System.out.println("WebSocket连接已关闭: " + reason);
        closeLatch.countDown();
        return Listener.super.onClose(webSocket, statusCode, reason);
    }

    @Override
    public void onError(WebSocket webSocket, Throwable error) {
        System.err.println("WebSocket发生错误: " + error.getMessage());
        error.printStackTrace();
        closeLatch.countDown();
    }

    /**
     * 获取最终结果
     */
    public String getResult() throws InterruptedException {
        // 等待完成或超时(30秒)
        boolean completed = closeLatch.await(30, TimeUnit.SECONDS);
        if (!completed) {
            throw new RuntimeException("WebSocket等待响应超时");
        }
        return resultBuilder.toString();
    }
}
